# coding: utf-8

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from dataclasses import dataclass
import math

torch.manual_seed(1024)

@dataclass
class GPTConfig:
    block_size: int = 512  # 文本的最大长度, 或者叫 max_seq_len
    batch_size: int = 32
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768  # hidden_size, hidden_dim; 这里是 embedding_size
    hidden_dim: int = n_embd
    # 为了可以til_embedding_weight
    dropout: float = 0.1
    head_size: int = n_embd // n_head
    vocab_size: int = 50257  # 词表大小，这里是 gpt2 的 vocab_size
    
# 定义single head attention
class SingleHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.key = nn.Linear(config.hidden_dim, config.head_size)
        self.value = nn.Linear(config.hidden_dim, config.head_size)
        self.query = nn.Linear(config.hidden_dim, config.head_size)
        self.head_size = config.head_size
        
        # 新的attention mask 写法，通过 register_buffer 注册到模型中, 这样可以直接调用 self.attention_mask, 可以不用计算梯度， 节约内存和显存，提升速度
        # 这里的 attention_mask 是一个 block_size x block_size 的 下三角矩阵，用于 mask 掉 padding 部分
        self.register_buffer(
            "attention_mask", 
            # tril 矩阵, 下三角矩阵, 用于 mask 掉 padding 部分
            torch.tril(
                torch.ones(config.block_size, config.block_size)
                )
        )
        self.dropout = nn.Dropout(config.dropout)
        
    def forward(self, x):
        batch_size, seq_len, hidden_dim = x.shape
        k = self.key(x)
        q = self.query(x)
        v = self.value(x)
        weight = q @ k.transpose(-2, -1)  # @ 就是点乘， 是 torch.matmul
        weight = weight.masked_fill(
            self.attention_mask[:seq_len, :seq_len] == 0, float('-inf')
            )  # 这里的 masked_fill 就是把 attention_mask 矩阵中为 0 的地方，用 -inf (负无穷) 填充
        # 要注意计算weight的时候，要除以根号 d_k, 因为点乘之后，得到的结果是没有归一化的，需要除以根号 d_k 才是权重
        weight = F.softmax(weight, dim=-1) / math.sqrt(self.head_size)
        # 这里的 dropout 就是在 attention 之后，再加上一个 dropout，防止过拟合
        weight = self.dropout(weight)  
        output = weight @ v  # dropout 之后，再乘以 v 得到输出，不能乘以 v 后再 dropout，因为这会导致信息丢失
        return output
    
# 定于 multi head attention
class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        # 用 ModuleList 包装多个 SingleHeadAttention， 实现多头注意力（实际上就是简单的 single head concat）
        self.heads = nn.ModuleList(
            [SingleHeadAttention(config) for _ in range(config.n_head)]
            )
        
        # proj 层，用于将 multi head attention 的输出，映射到 hidden_dim 维度
        # 加上 之前的 k, v, q 层，整个模型的权重矩阵就有 4 个 （计算模型的可训练参数量）
        self.proj = nn.Linear(config.hidden_dim, config.hidden_dim)
        self.dropout = nn.Dropout(config.dropout)
    
    def forward(self, x):
        # 简单的 single head concat
        output = torch.cat(
            [h(x) for h in self.heads], dim=-1
        )
        output = self.proj(output)
        output = self.dropout(output)
        return output
    
# 定义 MLP 层（feed forward）, 或称 DNN 层
class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.hidden_dim, 4 * config.hidden_dim),  # 升维， 4 * hidden_dim 就是 MLP 层的输入维度, swiglu 只用升维到 8/3 倍
            nn.GELU(),
            nn.Linear(4 * config.hidden_dim, config.hidden_dim),  # 降维
            nn.Dropout(config.dropout)
        )
        
    def forward(self, x):
        return self.net(x)
    
# 定义 block 层
class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attn = MultiHeadAttention(config)
        self.mlp = FeedForward(config)
        # 定义 2 个 layer norm 层，用于对输入进行归一化，防止梯度爆炸或消失
        self.ln1 = nn.LayerNorm(config.hidden_dim)
        self.ln2 = nn.LayerNorm(config.hidden_dim)
        
    def forward(self, x):
        x = self.ln1(x)
        x = x + self.attn(x)
        x = self.ln2(x)
        x = x + self.mlp(x)
        return x
    
# 定义 GPT 模型
class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        # （embedding, position, norm, mlp, block) 五个部分
        # 现在的position, 由固定位置编码，变成 ROPE 旋转位置编码
        # 现在的norm，由 layer norm 变成 rms norm（Llama改进之——均方根层归一化RMSNorm）
        # 现在的mlp， 升级到了 swiglu 函数
        # mha -> gqa
        
        self.token_embedding_table = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding_table = nn.Embedding(config.block_size, config.n_embd)
        self.blocks = nn.Sequential(
            *[Block(config) for _ in range(config.n_layer)]
        )
        
        self.ln_final = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        
        # 非常重要, 现在的 SLM 模型, 会用 tie weight 来减少 embedding 参数量，增加主干网络参数的参数量
        # 如果 embedding 占的参数太大，那模型学习到的知识会更少
        
        # 这里的 tie weight 就是把 lm_head 的权重参数，直接赋值给 token_embedding_table 的权重参数，这样就实现了 embedding 参数的共享
        # 为什么 vocab_size * n_embd 可以等于 n_embd * vocab_size 呢, 因为这里的 embedding 是双向的，所以实际上是 n_embd * vocab_size
        self.token_embedding_table.weight = self.lm_head.weight
        
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            # 初始化为正态分布
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            # 初始化为正态分布
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            
    def forward(self, idx, targets=None):
        # idx 是 输入的 token_ids，targets 是 目标 token_ids, shape 要一样
        batch_size, seq_len = idx.size()  # (batch_size, seq_len)
        token_embd = self.token_embedding_table(idx)  # (batch_size, seq_len, n_embd)
        pos_emb = self.position_embedding_table(
            # 这里的 pos_emb 就是位置编码， 要确保位置编码和输入的idx在同一个device上
            torch.arange(seq_len, device=idx.device)
        )
        # 经典题目： token embedding 和 position embedding 如何结合？
        x = token_embd + pos_emb  # shape is (batch_size, seq_len, n_embd)
        x = self.blocks(x)
        x = self.ln_final(x)
        logits = self.lm_head(x)
        if targets is None:
            loss = None
        else:
            batch, seq_len, vocab_size = logits.size()
            logits = logits.view(batch * seq_len, vocab_size)
            targets = targets.view(batch * seq_len)
            loss = F.cross_entropy(logits, targets)
        return logits, loss
    
    def generate(self, idx, max_new_tokens):
        pass  # TODO